import os
import cv2
import random
import utils

SrcDir = r"F:\403\ImgOutput\Video2Jpg"
DstDir = r"F:\403\ImgOutput\NarrowPic"
RandLowerBound = 2
RandUpperBound = 10
Epochs = 10

NameNum = 0

for root, dirs, files in os.walk(SrcDir):
    for file in files:
        if file[-1] == 'l':
            for Epoch in range(Epochs):
                Img = cv2.imread(root + "\\" + file[:-3] + "jpg")
                Locations = utils.ReadXml(root + "\\" + file)
                ImgCp = Img.copy()
                for location in Locations:
                    RandExtraTop = random.randint(RandLowerBound, RandUpperBound)
                    RandExtraBottom = random.randint(RandLowerBound, RandUpperBound)
                    RandExtraLeft = random.randint(RandLowerBound, RandUpperBound)
                    RandExtraRight = random.randint(RandLowerBound, RandUpperBound)
                    ImgPadding = cv2.copyMakeBorder(ImgCp, RandExtraTop, RandExtraBottom, RandExtraLeft, RandExtraRight,
                                                    cv2.BORDER_CONSTANT, value=0)
                    LabelImg = ImgPadding[location[1]:location[3] + RandExtraTop + RandExtraBottom, location[0]:location[2] + RandExtraLeft + RandExtraRight]
                    # print("%d %d %d %d" % (RandExtraTop, RandExtraBottom, RandExtraLeft, RandExtraRight))
                    # cv2.imshow(file, LabelImg)
                    # cv2.waitKey(0)
                    # cv2.destroyAllWindows()
                    NameNum += 1
                    LabelLocations = [DstDir + "\\" + "%05d" % NameNum + ".xml"]
                    LabelLocation = [RandExtraLeft, RandExtraTop, RandExtraLeft + location[2] - location[0], RandExtraTop + location[3] - location[1], "narrow"]
                    LabelLocations.append(LabelLocation)
                    utils.WriteXml(LabelLocations, LabelImg.shape[1], LabelImg.shape[0])
                    cv2.imwrite(DstDir + "\\" + "%05d" % NameNum + ".jpg", LabelImg)
                    print("%05d  " % NameNum)
